//------------------------------------------------------------------------------
// Addin.cpp
//------------------------------------------------------------------------------
// File provided for Microsoft Shared Source.
// Copyright (c) Microsoft Corporation. All rights reserved.
//
// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF 
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO 
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A 
// PARTICULAR PURPOSE.
//------------------------------------------------------------------------------
//
// The VS Addin class
//
//------------------------------------------------------------------------------

#include "pch.h"
#include ".\\FlatSolutionExplorerUI\\resource.h"

// The following ext_ConnectMode value is not defined in the header, but I've seen references
// to it in documentation and other Add-In implementations
const int ext_cm_UISetup = 5;

PCWSTR c_szCommandName = L"FlatSolutionExplorer.Addin.ViewFlatSolutionExplorer";
PCWSTR c_szRegRoot = L"Software\\Flat Solution Explorer";

class ATL_NO_VTABLE CFseAddin : 
    public CComObjectRootEx<CComSingleThreadModel>,
    public CComCoClass<CFseAddin, &__uuidof(FseAddin)>,
    public IServiceProviderImpl<CFseAddin>,
    public IDispatchImpl<IDTCommandTarget, &__uuidof(IDTCommandTarget), &__uuidof(__EnvDTE), 8, 0>,
    public IDispatchImpl<IDTExtensibility2, &__uuidof(IDTExtensibility2), &__uuidof(__AddInDesignerObjects), 1, 0>
{
public:
    CFseAddin() {}

    BEGIN_REGISTRY_MAP(IDR_ADDIN)
        REGMAP_DYNAMIC_ENTRY("MODULEPATH", s_GetModulePath)
        REGMAP_STATIC_ENTRY("SATELLITENAME", "FlatSolutionExplorerUI.dll")
    END_REGISTRY_MAP()

    DECLARE_NOT_AGGREGATABLE(CFseAddin)
    
    BEGIN_COM_MAP(CFseAddin)
        COM_INTERFACE_ENTRY(IDTExtensibility2)
        COM_INTERFACE_ENTRY(IDTCommandTarget)
        COM_INTERFACE_ENTRY2(IDispatch, IDTExtensibility2)
        COM_INTERFACE_ENTRY(IServiceProvider)
    END_COM_MAP()

    BEGIN_SERVICE_MAP(CFseAddin)
        SERVICE_ENTRY_FWD(SID_SFseAddIn, _spvsDTE)
    END_SERVICE_MAP()

    DECLARE_PROTECT_FINAL_CONSTRUCT()

    // IDTExtensibility2
    IFACEMETHODIMP OnConnection(__in IDispatch *pdispApp, ext_ConnectMode ConnectMode, __in IDispatch *pdispAddin, __in SAFEARRAY **psaCustom);
    IFACEMETHODIMP OnDisconnection(ext_DisconnectMode RemoveMode, __in SAFEARRAY **ppsaCustom);
    IFACEMETHODIMP OnAddInsUpdate(__in SAFEARRAY **ppsaCustom);
    IFACEMETHODIMP OnStartupComplete(__in SAFEARRAY **ppsaCustom);
    IFACEMETHODIMP OnBeginShutdown(__in SAFEARRAY **ppsaCustom);

    // IDTCommandTarget
    IFACEMETHODIMP QueryStatus(__in BSTR bstrCmdName, vsCommandStatusTextWanted cstwNeededText, __out vsCommandStatus *pvscs, __out VARIANT *pvarCommandText);
    IFACEMETHODIMP Exec(__in BSTR bstrCmdName, vsCommandExecOption ExecuteOption, __in VARIANT *pvarIn, __out VARIANT *pvarOut, __out VARIANT_BOOL *pfHandled);

private:
    static HRESULT s_GetModulePath(CString &str);

    HRESULT _AddCommand();
    HRESULT _AddCommandToViewMenu(__in Command *pvsCmd);
    HRESULT _EnsureToolWindow();
    HRESULT _ShowFlatSolutionExplorer();

    class CSolutionEventSink :
        public CDTEEventSinkBase<1, CSolutionEventSink, _dispSolutionEvents>
    {
    public:
        CSolutionEventSink() {}
        ~CSolutionEventSink() {}

        BEGIN_SINK_MAP(CSolutionEventSink)
            SINK_ENTRY_EX(1, __uuidof(_dispSolutionEvents), 1, Opened)
            SINK_ENTRY_EX(1, __uuidof(_dispSolutionEvents), 3, AfterClosing)
        END_SINK_MAP()

        IFACEMETHODIMP Opened()
        {
            GET_EMBEDDED_PARENT(CFseAddin, x_snkSolution, pc);
            pc->_Opened();
            return S_OK;
        }

        IFACEMETHODIMP AfterClosing()
        {
            GET_EMBEDDED_PARENT(CFseAddin, x_snkSolution, pc);
            pc->_AfterClosing();
            return S_OK;
        }

    protected:
        virtual HRESULT GetEventSourceObject(__in Events2 *pvsEvents2, REFIID riid, __deref_out void **ppv)
        {
            CComPtr<_SolutionEvents> spvsSolutionEvents;
            HRESULT hr = pvsEvents2->get_SolutionEvents(&spvsSolutionEvents);
            if (SUCCEEDED(hr))
            {
                hr = spvsSolutionEvents->QueryInterface(riid, ppv);
            }
            return hr;
        }

    } x_snkSolution;

    void _Opened();
    void _AfterClosing();

private:
    CComPtr<DTE2> _spvsDTE;
    CComPtr<AddIn> _spvsAddin;
    CComPtr<Window> _spvsWnd;
    CComPtr<IDispatch> _spdispFlatSolutionExplorer;
    CAtlArray<DWORD> _rgRegisteredObjects;
};

OBJECT_ENTRY_AUTO(__uuidof(FseAddin), CFseAddin)

// CFseAddin

HRESULT CFseAddin::s_GetModulePath(CString &str)
{
    PWSTR psz = str.GetBufferSetLength(MAX_PATH);
    HRESULT hr = GetModuleFileName(_AtlBaseModule.GetModuleInstance(), psz, MAX_PATH) ? S_OK : AtlHresultFromLastError();
    if (SUCCEEDED(hr))
    {
        PathRemoveFileSpec(psz);
    }
    return hr;
}

HRESULT CFseAddin::_AddCommandToViewMenu(__in Command *pvsCmd)
{
    CComPtr<_CommandBars> spvsCommandBars;
    HRESULT hr = _spvsDTE->get_CommandBars((IDispatch **)&spvsCommandBars);
    if (SUCCEEDED(hr))
    {
        try
        {
            CComPtr<CommandBar> spvsCommandBarMenu;
            hr = spvsCommandBars->get_Item(CComVariant(L"MenuBar"), &spvsCommandBarMenu);
            if (SUCCEEDED(hr))
            {
                CComPtr<CommandBarControls> spvsCommandBarControlsMenuBar;
                hr = spvsCommandBarMenu->get_Controls(&spvsCommandBarControlsMenuBar);
                if (SUCCEEDED(hr))
                {
                    CComPtr<CommandBarControl> spvsCommandBarControlView;
                    hr = spvsCommandBarControlsMenuBar->get_Item(CComVariant(L"View"), &spvsCommandBarControlView);
                    if (SUCCEEDED(hr))
                    {
                        // For some reason, I can't seem to convert from CommandBarControl to
                        // CommandBarPopup via COM (QI on CommandBarControl for CommandBarPopup
                        // fails even though I know it's a popup and popup derives from control).  For
                        // some sanity, check the type and then force the cast.
                        MsoControlType msoct;
                        hr = spvsCommandBarControlView->get_Type(&msoct);
                        if (SUCCEEDED(hr) && msoct == msoControlPopup)
                        {
                            CComPtr<CommandBarPopup> spvsCommandBarPopup = (CommandBarPopup *)spvsCommandBarControlView.p;
                            CComPtr<CommandBar> spvsCommandBarView;
                            hr = spvsCommandBarPopup->get_CommandBar(&spvsCommandBarView);
                            if (SUCCEEDED(hr))
                            {
                                CComPtr<CommandBarControl> spvsCommandBarControl;
                                hr = pvsCmd->AddControl(spvsCommandBarView, 1, (IDispatch **)&spvsCommandBarControl);
                            }
                        }
                    }
                }
            }
        }
        catch (CAtlException &e)
        {
            hr = e.m_hr;
        }
    }
    return hr;
}

HRESULT CFseAddin::_AddCommand()
{
    CComPtr<Commands> spvsCmds;
    HRESULT hr = _spvsDTE->get_Commands(&spvsCmds);
    if (SUCCEEDED(hr))
    {
        // See if the command already exists
        CComBSTR sbstrCmdID;
        hr = sbstrCmdID.Append(c_szCommandName);
        if (SUCCEEDED(hr))
        {
            CComVariant svarCmdID;
            svarCmdID.bstrVal = sbstrCmdID.Detach();
            svarCmdID.vt = VT_BSTR;
            CComPtr<Command> spvsCmd;
            hr = spvsCmds->Item(svarCmdID, -1, &spvsCmd);
            if (FAILED(hr))
            {
                // Command is not present.  Add it to the pool of commands.
                CComPtr<Commands2> spvsCmds2;
                hr = spvsCmds->QueryInterface(&spvsCmds2);
                if (SUCCEEDED(hr))
                {
                    hr = sbstrCmdID.Append(L"ViewFlatSolutionExplorer");
                    if (SUCCEEDED(hr))
                    {
                        CComBSTR sbstrCmdName;
                        hr = sbstrCmdName.LoadString(IDS_COMMANDNAME) ? S_OK : E_OUTOFMEMORY;
                        if (SUCCEEDED(hr))
                        {
                            CComBSTR sbstrCmdDescription;
                            hr = sbstrCmdDescription.LoadString(IDS_COMMANDDESCRIPTION) ? S_OK : E_OUTOFMEMORY;
                            if (SUCCEEDED(hr))
                            {
                                hr = spvsCmds2->AddNamedCommand2(_spvsAddin, sbstrCmdID, sbstrCmdName, sbstrCmdDescription, VARIANT_FALSE, CComVariant(IDB_FSECMD),
                                                                 NULL, vsCommandStatusSupported+vsCommandStatusEnabled, vsCommandStylePictAndText,
                                                                 vsCommandControlTypeButton, &spvsCmd);
                                hr = (SUCCEEDED(hr) && spvsCmd) ? S_OK : E_FAIL;
                                if (SUCCEEDED(hr))
                                {
                                    CRegKey rk;
                                    hr = AtlHresultFromWin32(rk.Create(HKEY_CURRENT_USER, c_szRegRoot));
                                    if (SUCCEEDED(hr))
                                    {
                                        BOOL fAddCommandToViewMenu = FALSE;
                                        if (SUCCEEDED(AtlHresultFromWin32(rk.QueryDWORDValue(L"AddCommandToView", (DWORD &)fAddCommandToViewMenu))))
                                        {
                                            rk.DeleteValue(L"AddCommandToView");
                                        }
                                        if (fAddCommandToViewMenu)
                                        {
                                            hr = _AddCommandToViewMenu(spvsCmd);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    return hr;
}

STDMETHODIMP CFseAddin::OnConnection(__in IDispatch *pdispApp, ext_ConnectMode ConnectMode, __in IDispatch *pdispAddin, __in SAFEARRAY **ppsaCustom)
{
    HRESULT hr = pdispApp->QueryInterface(__uuidof(DTE2), (void **)&_spvsDTE);
    if (SUCCEEDED(hr))
    {
        hr = pdispAddin->QueryInterface(__uuidof(AddIn), (void **)&_spvsAddin);
        if (SUCCEEDED(hr))
        {
            switch (ConnectMode)
            {
            case ext_cm_UISetup:
                hr = _AddCommand();
                break;

            default:
                {
                    static const CLSID *s_rgUnregistered[] =
                    {
                        &__uuidof(NullCategorizer),
                        &__uuidof(HitTypeCategorizer),
                        &__uuidof(TypeCategorizer),
                    };

                    try
                    {
                        for (int i = 0; i < ARRAYSIZE(s_rgUnregistered) && SUCCEEDED(hr); i++)
                        {
                            CComPtr<IUnknown> spunkFactory;
                            hr = _AtlModule.GetClassObject(*s_rgUnregistered[i], IID_PPV_ARGS(&spunkFactory));
                            if (SUCCEEDED(hr))
                            {
                                DWORD dwRegister;
                                hr = CoRegisterClassObject(*s_rgUnregistered[i], spunkFactory, CLSCTX_INPROC_SERVER, REGCLS_MULTIPLEUSE, &dwRegister);
                                if (SUCCEEDED(hr))
                                {
                                    _rgRegisteredObjects.Add(dwRegister);
                                }
                            }
                        }
                    }
                    catch (CAtlException &e)
                    {
                        hr = e.m_hr;
                    }

                    if (SUCCEEDED(hr))
                    {
                        x_snkSolution.Connect(_spvsDTE);
                    }
                }
                break;

            }
        }
    }
    return hr;
}

STDMETHODIMP CFseAddin::OnDisconnection(ext_DisconnectMode /*RemoveMode*/, __in SAFEARRAY **ppsaCustom)
{
    if (_spdispFlatSolutionExplorer)
    {
        CComPtr<IObjectWithSite> spObjectWithSite;
        if (SUCCEEDED(_spdispFlatSolutionExplorer->QueryInterface(&spObjectWithSite)))
        {
            spObjectWithSite->SetSite(NULL);
        }
    }

    x_snkSolution.Disconnect();

    for (size_t i = 0; i < _rgRegisteredObjects.GetCount(); i++)
    {
        CoRevokeClassObject(_rgRegisteredObjects.GetAt(i));
    }

    _spvsDTE = NULL;
    _spvsAddin = NULL;
    _spdispFlatSolutionExplorer = NULL;
    _spvsWnd = NULL;
    return S_OK;
}

STDMETHODIMP CFseAddin::OnAddInsUpdate(__in SAFEARRAY **ppsaCustom)
{
    return S_OK;
}

STDMETHODIMP CFseAddin::OnStartupComplete(__in SAFEARRAY **ppsaCustom)
{
    CRegKey rk;
    LONG lret = rk.Open(HKEY_CURRENT_USER, c_szRegRoot);
    if (lret == ERROR_SUCCESS)
    {
        DWORD dwVisible = 0;
        rk.QueryDWORDValue(L"Visible", dwVisible);
        if (dwVisible)
        {
            _ShowFlatSolutionExplorer();
        }
    }
    return S_OK;
}

STDMETHODIMP CFseAddin::OnBeginShutdown(__in SAFEARRAY **ppsaCustom)
{
    if (_spvsWnd)
    {
        VARIANT_BOOL fVisible;
        HRESULT hr = _spvsWnd->get_Visible(&fVisible);
        if (SUCCEEDED(hr))
        {
            CRegKey rk;
            hr = AtlHresultFromWin32(rk.Open(HKEY_CURRENT_USER, c_szRegRoot));
            if (SUCCEEDED(hr))
            {
                rk.SetDWORDValue(L"Visible", !!fVisible);
            }
        }
        _spvsWnd->Close(vsSaveChangesYes);
    }
    return S_OK;
}

STDMETHODIMP CFseAddin::QueryStatus(__in BSTR bstrCmdName, vsCommandStatusTextWanted cstwNeededText, __out vsCommandStatus *pvscs, __out VARIANT *pvarCommandText)
{
    if (cstwNeededText == vsCommandStatusTextWantedNone)
    {
        if (StrCmpI(bstrCmdName, c_szCommandName) == 0)
        {
            *pvscs = (vsCommandStatus)(vsCommandStatusSupported + vsCommandStatusEnabled);
        }
    }
    return S_OK;
}

HRESULT CFseAddin::_EnsureToolWindow()
{
    HRESULT hr = S_OK;
    if (!_spvsWnd)
    {
        CComPtr<Windows> spvsWindows;
        hr = _spvsDTE->get_Windows(&spvsWindows);
        if (SUCCEEDED(hr))
        {
            CComBSTR sbstrProgID;
            hr = sbstrProgID.Append(L"FlatSolutionExplorer.Control");
            if (SUCCEEDED(hr))
            {
                CComBSTR sbstrTitle;
                hr = sbstrTitle.LoadString(IDS_FSETITLENOSOLUTION) ? S_OK : E_OUTOFMEMORY;
                if (SUCCEEDED(hr))
                {
                    CComBSTR sbstrPosition;
                    hr = sbstrPosition.Append(L"{D9AB8B7C-6FC5-4b53-8F53-4F600CDD9EBB}");
                    if (SUCCEEDED(hr))
                    {
                        hr = spvsWindows->CreateToolWindow(_spvsAddin, sbstrProgID, sbstrTitle, sbstrPosition, &_spdispFlatSolutionExplorer, &_spvsWnd);
                        if (SUCCEEDED(hr))
                        {
                            CComPtr<IObjectWithSite> spObjectWithSite;
                            hr = _spdispFlatSolutionExplorer->QueryInterface(&spObjectWithSite);
                            if (SUCCEEDED(hr))
                            {
                                CComPtr<IUnknown> spunkThis;
                                hr = QueryInterface(IID_PPV_ARGS(&spunkThis));
                                if (SUCCEEDED(hr))
                                {
                                    spObjectWithSite->SetSite(spunkThis);
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    return hr;
}

HRESULT CFseAddin::_ShowFlatSolutionExplorer()
{
    HRESULT hr = _EnsureToolWindow();
    if (SUCCEEDED(hr))
    {
        VARIANT_BOOL fVisible;
        hr = _spvsWnd->get_Visible(&fVisible);
        if (SUCCEEDED(hr))
        {
            if (fVisible == VARIANT_FALSE)
            {
                hr = _spvsWnd->put_Visible(VARIANT_TRUE);
            }
            else
            {
                _spvsWnd->SetFocus();
            }
        }
    }
    return hr;
}

STDMETHODIMP CFseAddin::Exec(__in BSTR bstrCmdName, vsCommandExecOption vsceo, __in VARIANT *pvarIn, __out VARIANT *pvarOut, __out VARIANT_BOOL *pfHandled)
{
    *pfHandled = VARIANT_FALSE;
    if (vsceo == vsCommandExecOptionDoDefault)
    {
        if (StrCmpI(bstrCmdName, c_szCommandName) == 0)
        {
            _ShowFlatSolutionExplorer();
            *pfHandled = VARIANT_TRUE;
        }
    }
    return S_OK;
}

void CFseAddin::_Opened()
{
    if (_spvsWnd)
    {
        CComPtr<_Solution> spvsSolution;
        if (SUCCEEDED(_spvsDTE->get_Solution(&spvsSolution)))
        {
            CComBSTR sbstrName;
            if (SUCCEEDED(spvsSolution->get_FullName(&sbstrName)))
            {
                if (wcslen(sbstrName))
                {
                    try
                    {
                        CString str;
                        if (str.LoadString(IDS_FSETITLEOPEN))
                        {
                            CString strCaption;
                            strCaption.Format(str, sbstrName);

                            CComBSTR sbstrCaption;
                            if (SUCCEEDED(sbstrCaption.Append(strCaption)))
                            {
                                _spvsWnd->put_Caption(strCaption.GetBuffer());
                            }
                        }
                    }
                    catch (...)
                    {
                        // Nothing interesting to be done here
                    }
                }
            }
        }
    }
}

void CFseAddin::_AfterClosing()
{
    if (_spvsWnd)
    {
        CComBSTR sbstrName;
        if (sbstrName.LoadString(IDS_FSETITLENOSOLUTION))
        {
            _spvsWnd->put_Caption(sbstrName);
        }
    }
}
